# Author : ZZH
# Date : 2025/5/20
from pathlib import Path
import os
import yaml         # 路径操作
import shutil       # Yaml 配置文件处理
import logging      # 日志记录
import random
import xml.etree.ElementTree as ET      # XML解析库
from sklearn.model_selection import train_test_split        # 数据集分割


# 配置参数
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
print(PROJECT_ROOT)
INPUT_ANNOTATIONS = os.path.join(PROJECT_ROOT, "data", "raw", "annotations")
INPUT_IMAGES = os.path.join(PROJECT_ROOT, "data", "raw", "images")
OUTPUT_BASE = os.path.join(PROJECT_ROOT, "data")
CLASSES = set()         # 自动收集类别


# 创建输出目录结构
for subset in ["train", "test", "val"]:
    for dtype in ["images", "labels"]:
        os.makedirs(os.path.join(OUTPUT_BASE, subset, dtype), exist_ok=True)


def convert_coordinates(size, box):
    """将XML坐标转换为YOLO归一化坐标"""
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    return [x*dw, y*dh, w*dw, h*dh]


def parse_xml(xml_path):
    """解析单个XML标注文件"""
    tree = ET.parse(xml_path)
    root = tree.getroot()

    size = root.find("size")
    # if size is None:
    #     logging.error(f"XML文件{xml_path}缺少<size>标签")
    #     return
    width = int(size.find("width").text)
    height = int(size.find("height").text)

    objects = []
    for obj in root.iter("object"):
        cls = obj.find("name").text.strip()
        CLASSES.add(cls)
        bndbox = obj.find("bndbox")
        coords = [
            float(bndbox.find("xmin").text),
            float(bndbox.find("xmax").text),
            float(bndbox.find("ymin").text),
            float(bndbox.find("ymax").text)
        ]
        yolo_coords = convert_coordinates((width, height), coords)
        objects.append((cls, yolo_coords))

    return objects


def split_dataset(files):
    """8:1:1分割数据集"""
    random.shuffle(files)
    n = len(files)
    train_end = int(0.8 * n)
    test_end = train_end + int(0.1 * n)
    return {
        "train": files[:train_end],
        "test": files[train_end:test_end],
        "val": files[test_end:]
    }


# 主处理流程
if __name__ == "__main__":
    run_coed = 0
    # 处理所有XML文件
    dataset = []
    for xml_file in os.listdir(INPUT_ANNOTATIONS):
        if not xml_file.endswith(".xml"):
            continue
        # 获取 xml文件名字
        base_name = os.path.splitext(xml_file)[0]
        # 获取对应jpg文件路径
        img_file = next((f for f in os.listdir(INPUT_IMAGES)
                         if f.startswith(base_name)), None)
        if not img_file:
            continue

        # 解析转换标注
        xml_path = os.path.join(INPUT_ANNOTATIONS, xml_file)
        print(xml_path)
        objects = parse_xml(xml_path)

        # 保存标签文件
        label_content = []
        for cls, coords in objects:
            cls_id = sorted(CLASSES).index(cls)
            label_content.append(f"{cls_id} {' '.join(map(str, coords))}")

        dataset.append({
            "base_name": base_name,
            "img_file": img_file,
            "label_content": "\n".join(label_content)
        })

    # 数据集分割
    split = split_dataset(dataset)

    # 复制文件到对应目录
    for subset, items in split.items():
        for item in items:
            # 复制图片
            src_img = os.path.join(INPUT_IMAGES, item["img_file"])
            dst_img = os.path.join(OUTPUT_BASE, subset, "images", item["img_file"])
            shutil.copy(src_img, dst_img)

            # 保存标签
            label_path = os.path.join(OUTPUT_BASE, subset, "labels",
                                      f"{item['base_name']}.txt")
            with open(label_path, "w") as f:
                f.write(item["label_content"])

    # 生成data.yaml
    config = {
        "train": r"../../data/train/images",
        "val": r"../../data/val/images",
        "test": r"../../data/test/images",
        "names": sorted(list(CLASSES)),
        "nc": len(CLASSES)
    }

    with open(r"../../config/data.yaml", "w") as f:
        yaml.dump(config, f, default_flow_style=False)
